from typing import Optional, List
import os
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate
from langchain_community.llms import QianfanLLMEndpoint

AK = '2jCwbiCWReSVlPbt53f9pYhP'
SK = 'dwJ8BULLCsBMXsmtjyxNI8hkw4y0HPXq'
os.environ['QIANFAN_AK'] = AK
os.environ['QIANFAN_SK'] = SK
examples = [
    {'word': '开心', 'antonym': '难过'},
    {'word': '高', 'antonym': '矮'}
]
example_template = '''
单词:{word}
反义词{antonym}\\n
'''
example_prompt = PromptTemplate(
    input_variables=['word', 'antonym'],
    template=example_template
)
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    prefix='给每个单词的反义词',
    suffix=f'单词:{{input}}\n反义词:',
    input_variables=['input'],
    example_separator='\\n'
)
prompt_text = few_shot_prompt.format(input='粗')
print(prompt_text)
print('*' * 80)
# print('prompt_text-->',prompt_text)
llm = QianfanLLMEndpoint()
result = llm(prompt_text)
print(result)
